# -*- encoding: utf-8 -*-

"""
@File    :   ernie_model.py
@Time    :   2024/5/12
@Author  :   Wang Tengfei 
@Version :   1.0
@Email   :   wangtfcyj@163.com.com
"""

import os
import sys
import json
import requests


ernie_map = {
    'ERNIE-4.0-8K': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=',
    'ERNIE-3.5-8K': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token='
}


def get_access_token():
    """
    使用 API Key，Secret Key 获取access_token，替换下列示例中的应用API Key、应用Secret Key
    """

    url = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=va97KRVfThHtF88DJinJyOo6&client_secret=VBTaSDU7OBpCYF7oPQyQCgUrZblPCWKe"

    payload = json.dumps("")
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)
    return response.json().get("access_token")


def ernie_process():
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + get_access_token()

    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": "你好"
            }
        ],
        "stream": True
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload, stream=True)

    return response


def completion_to_openai_stream_response(ernie_data, model='default'):
    completion_tokens = ernie_data.get('usage', {}).get("completion_tokens")
    prompt_tokens = ernie_data.get('usage', {}).get("prompt_tokens", 0)
    total_tokens = ernie_data.get('usage', {}).get("total_tokens", 0)
    id = ernie_data.get('id')
    if ernie_data.get('is_end'):
        finish_reason = 'stop'
    else:
        finish_reason = None
    created = ernie_data.get('created')
    index = ernie_data.get("index", 0)
    if not ernie_data.get('is_end'):
        usage = None
    else:
        usage = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens,
        }

    openai_response = {
        "id": id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": model,
        "usage": usage,
        "choices": [
            {
                "delta": {
                    "role": "assistant",
                    "content": ernie_data.get('result'),
                },
                "index": index,
                "finish_reason": finish_reason,
            }
        ],
    }
    return openai_response


def to_stream():
    response = ernie_process()
    for line in response.iter_lines():
        if line.decode("UTF-8") == '':
            continue
        res = json.loads(line.decode("UTF-8").strip('data:'))
        yield completion_to_openai_stream_response(res)


if __name__ == '__main__':
    for x in to_stream():
        print(x)